import PyPDF2
import re
import ollama
import os
import glob
import pandas as pd
import threading
import time


# 定义函数来去除特殊空格和非法字符
def clean_text(text):
    # 移除特定的非法字符
    text = re.sub(r'[^\x20-\x7E]+', '', text)  # 只保留可打印的 ASCII 字符
    # 替换多个空格
    return re.sub(r'\s+', ' ', text).strip()


# 定义函数来处理单个PDF文件
def process_pdf(pdf_path, output_path):
    """
    处理单个PDF文件，提取文本并输出到指定路径。
    """
    try:
        with open(pdf_path, "rb") as file:
            reader = PyPDF2.PdfReader(file)
            with open(output_path, "w", encoding='utf-8') as output_file:
                for page in reader.pages:
                    text = page.extract_text()
                    if text:  # 检查是否成功提取文本
                        clean_text_result = clean_text(text)  # 清理文本
                        output_file.write(clean_text_result + "\n")  # 写入文件
                    else:
                        output_file.write("未提取到有效文本\n")
    except FileNotFoundError:
        print(f"文件未找到: {pdf_path}")
        return False
    except PyPDF2.errors.PdfReadError:
        print(f"无法读取PDF文件: {pdf_path}")
        return False
    except Exception as e:
        print(f"处理PDF文件时发生错误: {pdf_path}, 错误信息: {e}")
        return False
    return True


# 定义超时处理异常类
class TimeoutException(Exception):
    pass


# 定义带超时功能的线程类
class TimeoutThread(threading.Thread):
    """
    允许超时处理的线程类。
    """

    def __init__(self, target, args=(), kwargs={}):
        threading.Thread.__init__(self)
        self.target = target
        self.args = args
        self.kwargs = kwargs
        self.result = None
        self.exception = None

    def run(self):
        try:
            self.result = self.target(*self.args, **self.kwargs)
        except Exception as e:
            self.exception = e

    def join(self, timeout=None):
        super(TimeoutThread, self).join(timeout)
        if self.is_alive():
            raise TimeoutException("处理超时")
        if self.exception:
            raise self.exception
        return self.result


# 定义函数来处理文件夹中的所有PDF文件
def process_folder(folder_path, output_folder, excel_path):
    """
    处理指定文件夹中的所有PDF文件，并将结果保存到Excel文件中。
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    pdf_files = glob.glob(os.path.join(folder_path, "*.pdf"))
    results = []
    total_files = len(pdf_files)
    processed_files = 0
    errors = []
    unprocessed_files = []

    for pdf_file in pdf_files:
        base_name = os.path.basename(pdf_file).replace(".pdf", ".txt")
        output_path = os.path.join(output_folder, base_name)
        success = process_pdf(pdf_file, output_path)

        if not success:
            errors.append(pdf_file)
            continue

        with open(output_path, "r", encoding='utf-8') as file:
            content = file.read()

        try:
            # 使用线程实现超时处理
            def process_model():
                title = base_name.split(".txt")[0]
                res = ollama.chat(model='qwen2.5:14b', stream=False,
                                  messages=[{"role": "user", "content": f"{content}总结成摘要和关键词"}],
                                  options={"temperature": 0})
                summary = res['message']['content'].split('### 摘要\n\n')[1].split('\n\n### 关键词')[0]
                keywords = res['message']['content'].split('### 关键词\n\n')[1].split('\n- ')[1:]
                keywords = '、'.join(keywords)
                results.append({"文件名": title, "摘要": summary, "关键词": keywords})
                print(res)

            timeout_thread = TimeoutThread(target=process_model)
            timeout_thread.start()
            timeout_thread.join(timeout=30)

        except TimeoutException:
            print(f"处理大模型时超时: {pdf_file}")
            errors.append(pdf_file)
        except Exception as e:
            print(f"处理大模型时发生错误: {pdf_file}, 错误信息: {e}")
            errors.append(pdf_file)

        processed_files += 1
        print(f"进度: {processed_files}/{total_files} 文件已处理")

        # 每次处理完一个文件后保存Excel文件
        write_to_excel(results, excel_path)

    # 记录未处理的文件
    unprocessed_files = pdf_files[processed_files:]

    return results, errors, unprocessed_files


# 定义函数来将结果写入Excel文件
def write_to_excel(results, excel_path):
    """
    将处理结果写入指定的Excel文件。
    """
    df = pd.DataFrame(results)
    df.to_excel(excel_path, index=False)


# 主程序
if __name__ == "__main__":
    a = input("PDF文件夹路径:")
    b = input("TXT文件输出路径：")
    c = input("EXCEl文件输出路径:")
    folder_path = fr"{a}"  # 文件夹路径
    output_folder = fr"{b}"  # TXT文件输出路径
    excel_path = fr"{c}\results.xlsx"  # Excel文件输出路径

    results, errors, unprocessed_files = process_folder(folder_path, output_folder, excel_path)
    print(f"所有PDF文件已处理完毕，结果已保存到 {excel_path}")
    if errors:
        print("以下PDF文件处理失败:")
        for error in errors:
            print(error)
    if unprocessed_files:
        print("以下PDF文件未处理:")
        for unprocessed in unprocessed_files:
            print(unprocessed)